{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "torch.manual_seed(0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "  # Data Generation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dataset:\n",
    "    def __init__(self, vector_size=6, seed=None):\n",
    "        self.rng = np.random.RandomState(seed)\n",
    "        self.vector_size = vector_size\n",
    "        self.a_start = 0\n",
    "        self.a_end = 2\n",
    "        self.b_start = 4\n",
    "        self.b_end = 6\n",
    "\n",
    "    def batch(self, batch_size=128):\n",
    "        v = self.rng.uniform(0, 1, size=(batch_size, self.vector_size))\n",
    "        a = np.sum(v[:, self.a_start:self.a_end], axis=1)\n",
    "        b = np.sum(v[:, self.b_start:self.b_end], axis=1)\n",
    "        t = a + b\n",
    "\n",
    "        return (torch.tensor(v, dtype=torch.float32), torch.tensor(t[:, np.newaxis], dtype=torch.float32))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "  # Model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GradientBanditLayer(torch.nn.Module):\n",
    "    \"\"\"Implements the Gumbel NAC (Neural Accumulator)\n",
    "\n",
    "    Arguments:\n",
    "        in_features: number of ingoing features\n",
    "        out_features: number of outgoing features\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, in_features, out_features, name=None):\n",
    "        super().__init__()\n",
    "        self.name = name\n",
    "        self.in_features = in_features\n",
    "        self.out_features = out_features\n",
    "        self.beta = 0.9\n",
    "        self.alpha = 1e-3\n",
    "\n",
    "        self.register_buffer('samples', torch.Tensor(out_features, in_features, 3))\n",
    "        self.register_buffer('target_weights', torch.tensor([1, -1, 0], dtype=torch.float32))\n",
    "        self.running_mean_loss = torch.nn.Parameter(torch.tensor(0, dtype=torch.float32), requires_grad=False)\n",
    "\n",
    "        self.W_hat = torch.nn.Parameter(torch.Tensor(out_features, in_features, 3), requires_grad=False)\n",
    "        self.register_parameter('bias', None)\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        # Initialize to zero, the source of randomness can come from the Gumbel sampling.\n",
    "        torch.nn.init.zeros_(self.W_hat)\n",
    "\n",
    "    def set_iteration(self, iteration):\n",
    "        self.iteration = iteration\n",
    "\n",
    "    def optimize(self, loss):\n",
    "        pi = torch.nn.functional.softmax(self.W_hat, dim=-1)\n",
    "        self.running_mean_loss.mul_(self.beta).add_(1 - self.beta, loss)\n",
    "        running_mean_loss_debias = self.running_mean_loss / (1 - self.beta**(self.iteration + 1))\n",
    "        self.W_hat.addcmul_(self.alpha, running_mean_loss_debias - loss, self.samples - pi)\n",
    "\n",
    "    def forward(self, input):\n",
    "        # NOTE: This samples a W for all the observations, one could also sample a W for each observation.\n",
    "        # Similar approch could be done for GumbleSoftmax.\n",
    "        log_pi = torch.nn.functional.log_softmax(self.W_hat, dim=-1)\n",
    "        if self.iteration % 1000 == 0:\n",
    "            print(f'{self.name}.pi')\n",
    "            print(torch.exp(log_pi) @ self.target_weights)\n",
    "\n",
    "        self.samples = torch.nn.functional.gumbel_softmax(log_pi.view(-1, 3), hard=True).view(log_pi.size())\n",
    "        W = self.samples @ self.target_weights\n",
    "\n",
    "        return torch.nn.functional.linear(input, W, self.bias)\n",
    "\n",
    "    def extra_repr(self):\n",
    "        return 'in_features={}, out_features={}'.format(\n",
    "            self.in_features, self.out_features\n",
    "        )\n",
    "\n",
    "class Network(torch.nn.Module):\n",
    "    def __init__(self, model_name, vector_size=6):\n",
    "        super().__init__()\n",
    "        self.model_name = model_name\n",
    "\n",
    "        if model_name == 'GradientBandit':\n",
    "            self.layer_1 = GradientBanditLayer(vector_size, 2, name='layer_1')\n",
    "            self.layer_2 = GradientBanditLayer(2, 1, name='layer_2')\n",
    "        elif model_name == 'linear':\n",
    "            self.layer_1 = torch.nn.Linear(vector_size, 2)\n",
    "            self.layer_2 = torch.nn.Linear(2, 1)\n",
    "        else:\n",
    "            raise NotImplemented(f'{model_name} is not implemented')\n",
    "\n",
    "    def reset_parameters(self):\n",
    "        self.layer_1.reset_parameters()\n",
    "        self.layer_2.reset_parameters()\n",
    "\n",
    "    def set_iteration(self, iteration):\n",
    "        self.layer_1.set_iteration(iteration)\n",
    "        self.layer_2.set_iteration(iteration)\n",
    "\n",
    "    def optimize(self, loss):\n",
    "        self.layer_1.optimize(loss)\n",
    "        self.layer_2.optimize(loss)\n",
    "\n",
    "    def forward(self, input):\n",
    "        z_1 = self.layer_1(input)\n",
    "        z_2 = self.layer_2(z_1)\n",
    "        return z_2\n",
    "\n",
    "    def extra_repr(self):\n",
    "        return 'vector_size={}'.format(\n",
    "            self.vector_size\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "  # Training\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = Dataset(vector_size=6, seed=0)\n",
    "model = Network('GradientBandit', vector_size=6)\n",
    "model.reset_parameters()\n",
    "\n",
    "criterion = torch.nn.MSELoss()\n",
    "optimizer = torch.optim.Adam(model.parameters())\n",
    "\n",
    "for epoch_i in range(0, 100000):\n",
    "    model.set_iteration(epoch_i)\n",
    "\n",
    "    # Prepear\n",
    "    x, t = dataset.batch()\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    # Loss\n",
    "    y = model(x)\n",
    "    loss = criterion(y, t)\n",
    "\n",
    "    # Optimize\n",
    "    #loss.backward()\n",
    "    model.optimize(loss)\n",
    "    #optimizer.step()\n",
    "\n",
    "    if epoch_i % 1000 == 0:\n",
    "        print(f'train {epoch_i}: {loss}')\n"
   ]
  }
 ],
 "metadata": {
  "file_extension": ".py",
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  },
  "mimetype": "text/x-python",
  "name": "python",
  "npconvert_exporter": "python",
  "pygments_lexer": "ipython3",
  "version": 3
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
